import torch.nn as nn
from torchvision.models.vgg import vgg16
from mmcv.cnn import constant_init, kaiming_init, normal_init, xavier_init

from ..module import BACKBONE


@BACKBONE.register_module
class Vgg16(nn.Module):
    def __init__(self, pretrained):
        super(Vgg16, self).__init__()
        self.pretrained = pretrained
        self.features = nn.Sequential(*list(vgg16(pretrained=pretrained).features)[:-1])

    def init_weights(self):
        if self.pretrained:
            pass
        else:
            for m in self.features.modules():
                if isinstance(m, nn.Conv2d):
                    kaiming_init(m)
                elif isinstance(m, nn.BatchNorm2d):
                    constant_init(m, 1)
                elif isinstance(m, nn.Linear):
                    normal_init(m, std=0.01)

    def forward(self, x):
        outs = []
        for i, layer in enumerate(self.features):
            x = layer(x)
        outs.append(x)
        if len(outs) > 1:
            return tuple(outs)
        return outs[0]